#!/usr/bin/env python

import math
import rospy
#import signal
import tf
import time

from geometry_msgs.msg import Quaternion, Twist
from nav_msgs.msg import Odometry
from segway_msgs.msg import ConfigCmd
from segway_msgs.msg import Status
from segway_msgs.msg import AuxPower
from sensor_msgs.msg import BatteryState

class Controller():

    def __init__(self):

        self.cmd_pub = rospy.Publisher('/segway/cmd_vel', Twist, queue_size=10)
        self.odom_pub = rospy.Publisher('odom', Odometry, queue_size=10)
        self.cfg_pub = rospy.Publisher('/segway/gp_command', ConfigCmd, queue_size=10)
	self.bat_pub = rospy.Publisher('/battery0', BatteryState, queue_size=10)

        self.cmd_sub = rospy.Subscriber('cmd_vel', Twist, self.cmd_callback)
        self.odom_sub = rospy.Subscriber('/segway/feedback/wheel_odometry', Odometry, self.odom_callback)
        self.status_sub = rospy.Subscriber('/segway/feedback/status', Status, self.status_callback)
	self.bat_sub = rospy.Subscriber('/segway/feedback/aux_power', AuxPower, self.battery_callback)

        self.odom_frame_id = rospy.get_param('~odom_frame_id', 'odom')
        self.robot_frame_id = rospy.get_param('~robot_frame_id', 'base_footprint')
        self.publish_tf = rospy.get_param('~publish_tf', True)

        self.stop_timer = None

        self.target_linear_vel = 0.0
        self.target_angular_vel = 0.0

        self.terminate = False

        self.frame_rate = rospy.get_param('~cmd_publish_rate', 20)

        # Supplied in m/s^2
        self.linear_pos_accel_limit = rospy.get_param('~linear_pos_accel_limit', 1.0)
        self.linear_neg_accel_limit = rospy.get_param('~linear_neg_accel_limit', 1.0)
        self.angular_pos_accel_limit = rospy.get_param('~angular_pos_accel_limit', 1.0)
        self.angular_neg_accel_limit = rospy.get_param('~angular_neg_accel_limit', 1.0)

        # Converted to cmd publish rate
        self.linear_pos_accel_limit /= self.frame_rate
        self.linear_neg_accel_limit /= self.frame_rate
        self.angular_pos_accel_limit /= self.frame_rate
        self.angular_neg_accel_limit /= self.frame_rate

        # signal.signal(signal.SIGINT, self.sigint_handler)

    def status_callback(self, msg):
        if msg.operational_state == 3:
            cfg_cmd = ConfigCmd()
            cfg_cmd.gp_cmd = 'GENERAL_PURPOSE_CMD_SET_OPERATIONAL_MODE'
            cfg_cmd.gp_param = 5 # TRACTOR_REQUEST
            self.cfg_pub.publish(cfg_cmd)
        self.status_sub.unregister()

    def odom_callback(self, msg):
        msg.header.frame_id = self.odom_frame_id
        msg.child_frame_id = self.robot_frame_id

        # Modify position appropriately.
        msg.pose.pose.position.x = msg.pose.pose.position.x # reverse x
        msg.pose.pose.position.y = msg.pose.pose.position.y # reverse y
        quaternion = (msg.pose.pose.orientation.x,
                      msg.pose.pose.orientation.y,
                      msg.pose.pose.orientation.z,
                      msg.pose.pose.orientation.w)
        euler = tf.transformations.euler_from_quaternion(quaternion)
        quaternion = tf.transformations.quaternion_from_euler(euler[0], euler[1], euler[2] + math.pi)
        msg.pose.pose.orientation = Quaternion(*quaternion)

        # Modify velocity appropriately.
        msg.twist.twist.linear.x = -msg.twist.twist.linear.x

        if self.publish_tf:
            br = tf.TransformBroadcaster()
            br.sendTransform((msg.pose.pose.position.x, msg.pose.pose.position.y, msg.pose.pose.position.z),
                             quaternion,
                             msg.header.stamp,
                             self.robot_frame_id,
                             self.odom_frame_id)
        self.odom_pub.publish(msg)

    def timer_callback(self, event):
        self.target_linear_vel = 0
        self.target_angular_vel = 0

    def cmd_callback(self, msg):
        self.target_linear_vel = -msg.linear.x
        self.target_angular_vel = msg.angular.z
        if self.stop_timer is not None:
           self.stop_timer.shutdown()
        self.stop_timer = rospy.Timer(rospy.Duration(0.5), self.timer_callback)

    def battery_callback(self, msg):
	bat_msg = BatteryState()
	bat_msg.voltage = msg.aux_voltage_V[0]
        bat_msg.current = float('NaN')
        bat_msg.charge = float('NaN')
        bat_msg.capacity = float('NaN')
        bat_msg.design_capacity = float('NaN')
        bat_msg.percentage = float('NaN')
        bat_msg.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_DISCHARGING
        bat_msg.present = 1
	bat_msg.header = msg.header
	self.bat_pub.publish(bat_msg)

    def sigint_handler(self, signal, frame):
        self.terminate = True

    def run(self):
        rate = rospy.Rate(self.frame_rate)

        linear_vel = 0
        angular_vel = 0

        while (not rospy.is_shutdown()) and (not self.terminate):

            if linear_vel < self.target_linear_vel:
                # Increase speed
                if (self.linear_pos_accel_limit == 0.0 or
                    self.target_linear_vel - linear_vel < self.linear_pos_accel_limit):
                    linear_vel = self.target_linear_vel
                else:
                    linear_vel += self.linear_pos_accel_limit
            elif linear_vel > self.target_linear_vel:
                # Decrease speed
                if (self.linear_neg_accel_limit == 0.0 or
                    linear_vel - self.target_linear_vel < self.linear_neg_accel_limit):
                    linear_vel = self.target_linear_vel
                else:
                    linear_vel -= self.linear_neg_accel_limit

            if angular_vel < self.target_angular_vel:
                # Increase speed
                if (self.angular_pos_accel_limit == 0.0 or
                    self.target_angular_vel - angular_vel < self.angular_pos_accel_limit):
                    angular_vel = self.target_angular_vel
                else:
                    angular_vel += self.angular_pos_accel_limit
            elif angular_vel > self.target_angular_vel:
                # Decrease speed
                if (self.angular_neg_accel_limit == 0.0 or
                    angular_vel - self.target_angular_vel < self.angular_neg_accel_limit):
                    angular_vel = self.target_angular_vel
                else:
                    angular_vel -= self.angular_neg_accel_limit

            twist = Twist()
            twist.linear.x = linear_vel
            twist.angular.z = angular_vel
            self.cmd_pub.publish(twist)
            rate.sleep()

        # Shutdown base and switch to standby
        cfg_cmd = ConfigCmd()
        cfg_cmd.gp_cmd = 'GENERAL_PURPOSE_CMD_SET_OPERATIONAL_MODE'
        cfg_cmd.gp_param = 4 # STANDBY_REQUEST
        self.cfg_pub.publish(cfg_cmd)

if __name__=="__main__":

    rospy.init_node('segbot_v3_controller')
    controller = Controller()
    # TODO: figure out if segway is ready here, instead of arbitrary 5s sleep.
    time.sleep(5)
    controller.run()
